#Todo:
#sympy cse
#    if it is used only once, just assign it.
#printer subclass: numpy, pyopencl
#lambda vs function??
#move self._constant_ ....... to just _constant_
#[].append() vs [] += []      ???
#input error checking
#    proper constant size
#    duplicated variables?
#    proper function?
#    x and data checking
#    when it returns f,J, give an option to use preallocated j,J array
from __future__ import division
import sys
sys.path.insert(1, '/home/kkb110/research/workspace/sympy')
from pylab import *
import sympy
print sys.path
print sympy
from scipy import optimize
import pyopencl
from OpenclPrinter import OpenclPrinter
import itertools
import bisect
class fJ_generator:
    def __init__(self, function, variables, constants):
        variables = list(variables)
        constants = dict(constants)
        self._n = len(variables)
        self._m = 1
        function = function
        d = vars(sys.modules['numpy'])
        d.update([('pyopencl', pyopencl)])
        for i, ii in enumerate(variables):
            variables[i] = sympy.Symbol('_variable_' + str(ii))
            function = function.subs(ii, variables[i])

        for k, v in constants.items():
            constants[sympy.Symbol('_constant_' + str(k))] = array(v, dtype=float32, copy=True).reshape(-1, 1)
            del constants[k]
            #do we need the following line?
            #setattr(self, '_constant_' + str(k), array(v, copy=True).reshape(-1, 1))
            function = function.subs(k, sympy.Symbol('_constant_' + str(k)))
            self._m = max(self._m, len(v))

        d.update([(str(k), v) for (k, v) in constants.items()])

        jeq = []
        for i in variables:
            jeq.append(function.diff(i))

        #######################################################################
        class equation:
            def __init__(self, symbol=None, expression=None):
                self.symbol = symbol
                self.expression = expression
            def __repr__(self):
                return str(vars(self))
        test = [equation(*i) for i in sympy.cse([function] + jeq)[0]]

        constants2 = dict(constants)
        variables2 = list(variables)
        f_symbols = sympy.cse([function] + jeq)[1][0].atoms()
        f_expressions = []
        J_symbols = set([j for i in xrange(len(sympy.cse([function] + jeq)[1]))
                                for j in list(sympy.cse([function] + jeq)[1][i].atoms())])
        j_expressions = []

        for i in test:
            if i.expression.has(*variables2):
                i.has_variable = True
                variables2.append(i.symbol)
            else:
                i.has_variable = False
                constants2[i.symbol] = eval(str(i.expression), d)
                exec (str(i.symbol) + ' = ' + str(i.expression)) in d

        for i in reversed(test):
            if i.symbol.has(*J_symbols):
                i.used_in_J = True
                J_symbols.update(i.expression.atoms())
            else:
                i.used_in_J = False
                J_symbols.update(i.expression.atoms())
            if i.symbol.has(*f_symbols):
                i.used_in_f = True
                f_symbols.update(i.expression.atoms())
            else:
                i.used_in_f = False

        for i in test:
            if i.has_variable:
                if i.used_in_f:
                    f_expressions += [i]
                if i.used_in_J:
                    j_expressions += [i]

        print ('''
def _f(variables):
    variables = asarray(variables, dtype=float32).reshape(-1,{len_vars}).T
    ''' + str(variables) + ''' = [variables.take([i], 0) for i in xrange({len_vars})]
    ''' + '''
    '''.join([str(i.symbol) + ' = ' + str(i.expression) for i in f_expressions]) + '''
    return sum(''' + str(sympy.cse([function] + jeq)[1][0]) + ''',1)''').format(len_vars=len(variables))
    
        exec ('''
def _f(variables):
    variables = asarray(variables, dtype=float32).reshape(-1,{len_vars}).T
    ''' + str(variables) + ''' = [variables.take([i], 0) for i in xrange({len_vars})]
    ''' + '''
    '''.join([str(i.symbol) + ' = ' + str(i.expression) for i in f_expressions]) + '''
    return sum(''' + str(sympy.cse([function] + jeq)[1][0]) + ''',1)''').format(len_vars=len(variables)) in d
        setattr(self, '_f', d['_f'])

        exec ('''
def _J(variables):
    variables = asarray(variables, dtype=float32).reshape(-1,{len_vars}).T
    ''' + str(variables) + ''' = [variables.take([i], 0) for i in xrange({len_vars})]
    j = empty([''' + str(self._m) + ''',variables.size], dtype=float32)
    ''' + '''
    '''.join([str(i.symbol) + ' = ' + str(i.expression) for i in j_expressions]) + '''
    ''' + '''
    '''.join(['''j[:,{i}::{len_vars}]  =  '''.format(i=i, len_vars=len(variables)) + str(sympy.cse([function] + jeq)[1][i + 1]) for i in xrange(len(variables))]) + '''
    return j''').format(len_vars=len(variables)) in d
        setattr(self, '_J', d['_J'])

        print ('''
def _J(variables):
    variables = asarray(variables, dtype=float32).reshape(-1,{len_vars}).T
    ''' + str(variables) + ''' = [variables.take([i], 0) for i in xrange({len_vars})]
    j = empty([''' + str(self._m) + ''',variables.size], dtype=float32)
    ''' + '''
    '''.join([str(i.symbol) + ' = ' + str(i.expression) for i in j_expressions]) + '''
    ''' + '''
    '''.join(['''j[:,{i}::{len_vars}]  =  '''.format(i=i, len_vars=len(variables)) + str(sympy.cse([function] + jeq)[1][i + 1]) for i in xrange(len(variables))]) + '''
    return j''').format(len_vars=len(variables))


        print dict(zip(variables, itertools.repeat('gid0')) + zip(constants2.keys(), itertools.repeat('gid1')))
        ############################################################ opencl part
        openclPrinter = OpenclPrinter({'symbol_default_index':dict(zip(variables, itertools.repeat('gid1')) + zip(constants2.keys(), itertools.repeat('gid0')))})

        opencl_f_code = '''
__kernel void f(''' + ', '.join(['__global const float *' + str(i) for i in variables + constants2.keys()] + ['__global float *result']) + ''')
{
    int gid0 = get_global_id(0);
    int gid1 = get_global_id(1);
    ''' + '''
    '''.join(['float ' + str(i.symbol) + ' = ' + openclPrinter.doprint(i.expression) + ';' for i in f_expressions]) + '''
    return ''' + openclPrinter.doprint(sympy.cse([function] + jeq)[1][0]) + ''';
}'''

        opencl_J_code = '''
__kernel void J(''' + ', '.join(['__global const float *' + str(i) for i in variables + constants2.keys()] + ['__global float *result']) + ''')
{
    int gid0 = get_global_id(0);
    int gid1 = get_global_id(1);
    ''' + '''
    '''.join(['float ' + str(i.symbol) + ' = ' + openclPrinter.doprint(i.expression) + ';' for i in j_expressions]) + '''
    __global float *j = result + gid1*''' + str(len(variables)) + ''' + gid0*''' + str(len(variables)) + '''*get_global_size(1);
    ''' + '''
    '''.join(['''j[{i}]  =  '''.format(i=i, len_vars=len(variables)) + openclPrinter.doprint(sympy.cse([function] + jeq)[1][i + 1]) + ';' for i in xrange(len(variables))]) + '''
}'''

        self._opencl_ctx = pyopencl.create_some_context()
        self._opencl_queue = pyopencl.CommandQueue(self._opencl_ctx)
        self._opencl_constants = [pyopencl.Buffer(self._opencl_ctx, pyopencl.mem_flags.READ_ONLY | pyopencl.mem_flags.COPY_HOST_PTR, hostbuf=i) for i in constants2.values()]

#        self._opencl_constants = {}
#        for k, v in constants2.items():
#            self._opencl_constants[k] = pyopencl.Buffer(self._opencl_ctx, pyopencl.mem_flags.READ_ONLY | pyopencl.mem_flags.COPY_HOST_PTR, hostbuf=v)

#        exec ('''
#def _opencl_f(variables):
#    variables = asarray(variables).reshape(-1,{len_vars}).T
#    ''' + str(variables) + ''' = [variables.take([i], 0) for i in xrange({len_vars})]
#    ''' + '''
#    return sum(''' + str(sympy.cse([function] + jeq)[1][0]) + ''',1)''').format(len_vars=len(variables)) in d
#        setattr(self, '_f', d['_f'])

        exec ('''
def _opencl_J(self, variables):
    variables = asarray(variables ,dtype=float32).reshape(-1,{len_vars}).T
    ''' + str(variables) + ''' = [variables.take([i], 0) for i in xrange({len_vars})]
    ''' + str(variables) + ''' = [pyopencl.Buffer(self._opencl_ctx, pyopencl.mem_flags.READ_ONLY | pyopencl.mem_flags.COPY_HOST_PTR, hostbuf=i) for i in ''' + str(variables) + ''']
    j = empty([''' + str(self._m) + ''',variables.size], dtype=float32)
    result = pyopencl.Buffer(self._opencl_ctx, pyopencl.mem_flags.WRITE_ONLY, j.nbytes)
    self._opencl_kernel_J.J(self._opencl_queue, (j.shape[0]+1,j.shape[1]//{len_vars}), None, *(''' + str(variables) + ''' + self._opencl_constants ''' + '''+ [result]))
    pyopencl.enqueue_read_buffer(self._opencl_queue, result, j).wait()
#    result.release()
#    for i in ''' + str(variables) + ''':
#        i.release()
    return j
    ''').format(len_vars=len(variables)) in d
        setattr(self, '_opencl_J', d['_opencl_J'])

#        dest_buf = pyopencl.Buffer(self._opencl_ctx, pyopencl.mem_flags.WRITE_ONLY, b.nbytes)
#        print opencl_f_code
        print opencl_J_code
        self._opencl_kernel_f = pyopencl.Program(self._opencl_ctx, opencl_f_code).build()
        self._opencl_kernel_J = pyopencl.Program(self._opencl_ctx, opencl_J_code).build()

    def f(self, x, data=None):
        if data == None:
            return self._f(x)
        else:
            return self._f(x) - asarray(data).ravel()

    def J(self, x, data=None):
        #return self._opencl_J(self, x)
        return self._J(x)

    def init_opencl(self,):
        pass




















if __name__ == '__main__':
    x0, x1, x2, x3, x4, x5 = sympy.symbols(['x0', 'x1', 'x2', 'x3', 'x4', 'x5'])
    t = sympy.Symbol('t')

    thecls = fJ_generator((x0 + x1 * t + x2 * t ** 2) * sympy.cos((x3 + x4 * t + x5 * t ** 2)), [x0, x1, x2, x3, x4, x5], {t:linspace(0, 1, 4410, endpoint=False)})
    data = thecls.f([20, -20, 1, 0, 20, 1])
    x = [3, 0, 3, 0, 25, -3]
    x = [1, 2, 3, 4, 5, 6]
    import time

    thecls.f(x, data)
    
    print thecls.J(x+x, data)
    
#    time.sleep(30)
    (result, success, r1, r2, r3) = optimize.leastsq(thecls.f, x, args=(data), Dfun=thecls.J, full_output=True)

    plot(thecls.f(result, zeros(4410)))
    plot(data, 'r--')
    plot(thecls.f(array(x), zeros(4410)), 'go')
    show()




